from os import PathLike
from pathlib import Path
from typing import AnyStr, Dict

import plyvel

from storage.kvstore.base import AbstractKVStore


class PlyvelKVStore(AbstractKVStore):
    def __init__(self, path: AnyStr | PathLike):
        self._path = Path(path)
        self._db = plyvel.DB(str(self._path), create_if_missing=True)

    def close(self):
        self._db.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get(self, k: bytes, default: bytes = None) -> bytes:
        return self._db.get(k, default)

    def put(self, k: bytes, v: bytes):
        return self._db.put(k, v)

    def update(self, kv_pairs: Dict[bytes, bytes]):
        with self._db.write_batch() as wb:
            for k, v in kv_pairs.items():
                wb.put(k, v)

    def __delitem__(self, k: bytes):
        self._db.delete(k)

    def __getitem__(self, item: bytes):
        if ret := self._db.get(item):
            return ret
        raise KeyError(item)

    def __setitem__(self, key: bytes, value: bytes):
        self._db.put(key, value)
